package cn.com.daimajiangxin.flink;

import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;

import java.time.Duration;

import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.util.Collector;

/**
 * 使用 KeyedProcessFunction 实现“伪窗口”。
 * 将事件按窗口大小（如 1 小时）分桶累加，并在窗口结束时通过定时器输出。
 * <p>
 * 输入：TaxiFare（包含 driverId、eventTime、tip）
 * 输出：Tuple3<driverId, windowStart, totalTip>
 */
public class PseudoWindow extends KeyedProcessFunction<Long, TaxiFare, Tuple3<Long, Long, Float>> {
    public static final long serialVersionUID = 1L;
    private final long windowSizeMillis;
    private transient MapState<Long, Float> sumByWindowStart; // windowStart -> sumTip

    public PseudoWindow(Duration windowSize) {
        this.windowSizeMillis = windowSize.toMillis();
    }

    @Override
    public void open(org.apache.flink.configuration.Configuration parameters) throws Exception {
        MapStateDescriptor<Long, Float> desc = new MapStateDescriptor<>(
                "sumByWindowStart",
                Long.class,
                Float.class);
        sumByWindowStart = getRuntimeContext().getMapState(desc);
    }

    @Override
    public void processElement(TaxiFare value,
                               KeyedProcessFunction<Long, TaxiFare, Tuple3<Long, Long, Float>>.Context ctx,
                               Collector<Tuple3<Long, Long, Float>> out) throws Exception {
        if (value == null)
            return;
        long eventTs = value.getEventTime();
        long windowStart = alignToWindowStart(eventTs);
        long windowEnd = windowStart + windowSizeMillis;

        // 累加该窗口内的小费
        Float current = sumByWindowStart.get(windowStart);
        if (current == null)
            current = 0f;
        sumByWindowStart.put(windowStart, current + value.getTip());

        // 为窗口结束注册 EventTime 定时器（重复注册由 Flink 去重）
        ctx.timerService().registerEventTimeTimer(windowEnd);
    }

    @Override
    public void onTimer(long timestamp,
                        KeyedProcessFunction<Long, TaxiFare, Tuple3<Long, Long, Float>>.OnTimerContext ctx,
                        Collector<Tuple3<Long, Long, Float>> out) throws Exception {
        long windowEnd = timestamp;
        long windowStart = windowEnd - windowSizeMillis;
        Float sum = sumByWindowStart.get(windowStart);
        if (sum != null) {
            long driverId = ctx.getCurrentKey();
            out.collect(Tuple3.of(driverId, windowStart, sum));
            // 清理该窗口的累加值，避免状态膨胀
            sumByWindowStart.remove(windowStart);
        }
    }

    private long alignToWindowStart(long ts) {
        return (ts / windowSizeMillis) * windowSizeMillis;
    }
}
